import pybuda._C
import pybuda._C.graph
import pybuda._C.scheduler
import typing
from _typeshed import Incomplete
from typing import ClassVar, Dict, List, Optional, Tuple, Union, overload

C: TStreamDir
CNN: PolicyType
CZ: TStreamDir
ConsumerOperandDataEdgesFirst: GraphSolverSelfCutType
FastCut: GraphSolverSelfCutType
MMIO_LAST: ChipPlacementPolicy
MinimizeGrid: PolicyType
NLP: PolicyType
ProducerUserDataEdgesFirst: GraphSolverSelfCutType
R: TStreamDir
RZ: TStreamDir
Random: PolicyType
Ribbon: PolicyType
SNAKE: ChipPlacementPolicy

class BalancerConfig:
    chip_ids: List[int]
    default_dram_parameters: bool
    device_config: Incomplete
    enable_auto_transposing_placement: bool
    enable_enumerate_u_kt: bool
    enable_single_buffer_fallback: bool
    enable_t_streaming: bool
    fork_join_tiles_treshold: int
    graph_solver_self_cut_type: GraphSolverSelfCutType
    input_queues_on_host: bool
    manual_t_streaming: bool
    op_name_to_placer_overrides: Dict[str, pybuda._C.placer.OpOverride]
    op_names_to_chip_break: List[List[str]]
    op_names_to_chip_id_assignment: List[List[str]]
    op_names_to_epoch_break: List[List[str]]
    op_overrides: Dict[str, OpOverride]
    output_queues_on_host: bool
    policy_type: PolicyType
    random_policy_seed: int
    scheduler_config: pybuda._C.scheduler.SchedulerConfig
    skip_l1_usage_validation: bool
    target_cycles_offset: int
    use_interactive_placer: bool
    def __init__(self, device_config, scheduler_config: pybuda._C.scheduler.SchedulerConfig, policy_type: PolicyType = ..., random_policy_seed: int = ..., chip_ids: List[int] = ..., chip_placement_policy: ChipPlacementPolicy = ..., default_dram_parameters: bool = ..., skip_l1_usage_validation: bool = ..., enable_t_streaming: bool = ..., manual_t_streaming: bool = ..., input_queues_on_host: bool = ..., output_queues_on_host: bool = ..., op_overrides: Dict[str, OpOverride] = ..., op_names_to_epoch_break: List[List[str]] = ..., op_names_to_chip_break: List[List[str]] = ..., op_names_to_chip_id_assignment: Dict[str, int] = ..., op_name_to_placer_overrides: Dict[str, pybuda._C.placer.OpOverride] = ..., enable_auto_transposing_placement: bool = ..., graph_solver_self_cut_type: GraphSolverSelfCutType = ..., use_interactive_placer: bool = ..., enable_enumerate_u_kt: bool = ..., enable_single_buffer_fallback: bool = ...) -> None: ...

class BalancerSolution:
    def __init__(self, *args, **kwargs) -> None: ...
    def cut_edges_as_override(self, arg0: pybuda._C.graph.Graph) -> List[Tuple[str, str, int]]: ...
    @property
    def op_models(self): ...
    @property
    def output_host_tms(self): ...
    @property
    def placer_solution(self) -> pybuda._C.placer.PlacerSolution: ...

class BlockShape:
    def __init__(self, *args, **kwargs) -> None: ...
    def buffer_tiles(self, arg0: int) -> int: ...
    def volume(self) -> int: ...
    def __eq__(self, arg0: BlockShape) -> bool: ...
    @property
    def mblock_m(self) -> int: ...
    @property
    def mblock_n(self) -> int: ...
    @property
    def t(self) -> int: ...
    @property
    def ublock(self) -> UBlockShape: ...

class BufferModel:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def block_shape(self) -> BlockShape: ...
    @property
    def data_format(self) -> pybuda._C.DataFormat: ...
    @property
    def l1_size_tiles(self) -> int: ...

class ChipPlacementPolicy:
    __members__: ClassVar[dict] = ...  # read-only
    MMIO_LAST: ClassVar[ChipPlacementPolicy] = ...
    SNAKE: ClassVar[ChipPlacementPolicy] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class Constant:
    def __init__(self, value: int) -> None: ...

class FactorizedInt:
    @overload
    def __init__(self) -> None: ...
    @overload
    def __init__(self, max_val: int) -> None: ...
    @overload
    def __init__(self, range: Tuple[int, int]) -> None: ...
    @overload
    def __init__(self, constant: Constant) -> None: ...
    def __and__(self, arg0: FactorizedInt) -> FactorizedInt: ...
    def __div__(self, arg0: FactorizedInt) -> FactorizedInt: ...
    def __mul__(self, arg0: FactorizedInt) -> FactorizedInt: ...
    def __or__(self, arg0: FactorizedInt) -> FactorizedInt: ...
    def __sub__(self, arg0: FactorizedInt) -> FactorizedInt: ...
    @property
    def factors(self) -> List[int]: ...
    @property
    def max_factor(self) -> int: ...

class FusedSubOpModel:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def has_dest_input(self) -> bool: ...
    @property
    def has_dest_output(self) -> bool: ...
    @property
    def mblock_k(self) -> int: ...
    @property
    def mblock_m(self) -> int: ...
    @property
    def mblock_n(self) -> int: ...
    @property
    def reduce_dim(self) -> str: ...
    @property
    def type(self) -> str: ...
    @property
    def ublock_ct(self) -> int: ...
    @property
    def ublock_kt(self) -> int: ...
    @property
    def ublock_rt(self) -> int: ...

class GraphSolverSelfCutType:
    __members__: ClassVar[dict] = ...  # read-only
    ConsumerOperandDataEdgesFirst: ClassVar[GraphSolverSelfCutType] = ...
    FastCut: ClassVar[GraphSolverSelfCutType] = ...
    ProducerUserDataEdgesFirst: ClassVar[GraphSolverSelfCutType] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class GridShape:
    def __init__(self, *args, **kwargs) -> None: ...
    @overload
    def __eq__(self, arg0: GridShape) -> bool: ...
    @overload
    def __eq__(self, arg0: Tuple[int, int]) -> bool: ...
    @property
    def c(self) -> int: ...
    @property
    def r(self) -> int: ...

class OpModel:
    def __init__(self, *args, **kwargs) -> None: ...
    def block_shape(self) -> BlockShape: ...
    def buda_op_attrs(self, *args, **kwargs): ...
    def get_reduce_dim(self) -> str: ...
    def math_fidelity(self, *args, **kwargs): ...
    def op_type(self) -> str: ...
    @property
    def data_format(self) -> pybuda._C.DataFormat: ...
    @property
    def grid_shape(self) -> GridShape: ...
    @property
    def input_buffers(self) -> List[BufferModel]: ...
    @property
    def is_sparse_matmul(self) -> bool: ...
    @property
    def nz_strips(self) -> int: ...
    @property
    def nz_tiles(self) -> int: ...
    @property
    def nz_ublocks(self) -> int: ...
    @property
    def op_shape(self) -> OpShape: ...
    @property
    def output_buffers(self) -> List[BufferModel]: ...
    @property
    def parameter_buffers(self) -> List[BufferModel]: ...
    @property
    def sparse_indices(self) -> int: ...
    @property
    def t_stream_factor(self) -> TStreamFactor: ...

class OpOverride:
    force_dram_parameters: Optional[bool]
    grid_shape: Optional[Tuple[int, int]]
    input_buffer_multiplier: Optional[Dict[int, int]]
    output_buffer_multiplier: Optional[int]
    t_stream_dir: str
    t_stream_shape: Optional[Tuple[int, int]]
    u_kt: Optional[int]
    def __init__(self) -> None: ...
    def from_json(self) -> OpOverride: ...
    def to_json(self) -> Dict[str, Union[Optional[bool], str, Optional[int], Optional[List[int[2]]], Optional[Dict[int, int]]]]: ...

class OpShape:
    def __init__(self, input_shapes: List[Tuple[int, int, int, int]], output_shape: Tuple[int, int, int, int], scalar_dims: bool = ...) -> None: ...
    @property
    def inputs(self) -> List[TensorShape]: ...
    @property
    def outputs(self) -> List[TensorShape]: ...

class OutputHostTM:
    def __init__(self, *args, **kwargs) -> None: ...
    @classmethod
    def from_json(cls, arg0: Dict[str, Union[bool, int]]) -> OutputHostTM: ...
    def to_json(self) -> Dict[str, Union[bool, int]]: ...
    @property
    def hstack_factor(self) -> int: ...
    @property
    def row_major(self) -> bool: ...
    @property
    def vstack_factor(self) -> int: ...

class PolicyType:
    __members__: ClassVar[dict] = ...  # read-only
    CNN: ClassVar[PolicyType] = ...
    MinimizeGrid: ClassVar[PolicyType] = ...
    NLP: ClassVar[PolicyType] = ...
    Random: ClassVar[PolicyType] = ...
    Ribbon: ClassVar[PolicyType] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class TStreamDir:
    __members__: ClassVar[dict] = ...  # read-only
    C: ClassVar[TStreamDir] = ...
    CZ: ClassVar[TStreamDir] = ...
    R: ClassVar[TStreamDir] = ...
    RZ: ClassVar[TStreamDir] = ...
    __entries: ClassVar[dict] = ...
    def __init__(self, value: int) -> None: ...
    def __eq__(self, other: object) -> bool: ...
    def __hash__(self) -> int: ...
    def __index__(self) -> int: ...
    def __int__(self) -> int: ...
    def __ne__(self, other: object) -> bool: ...
    @property
    def name(self) -> str: ...
    @property
    def value(self) -> int: ...

class TStreamFactor:
    def __init__(self, *args, **kwargs) -> None: ...
    def __eq__(self, arg0: TStreamFactor) -> bool: ...
    @property
    def c(self) -> int: ...
    @property
    def dir(self) -> TStreamDir: ...
    @property
    def r(self) -> int: ...

class TensorShape:
    def __init__(self, *args, **kwargs) -> None: ...
    def __getitem__(self, arg0: int) -> int: ...
    def __iter__(self) -> typing.Iterator[int]: ...
    def __setitem__(self, arg0: int, arg1: int) -> None: ...
    @property
    def ct(self) -> int: ...
    @property
    def rt(self) -> int: ...
    @property
    def w(self) -> int: ...
    @property
    def z(self) -> int: ...

class UBlockShape:
    def __init__(self, *args, **kwargs) -> None: ...
    def volume(self) -> int: ...
    @overload
    def __eq__(self, arg0: UBlockShape) -> bool: ...
    @overload
    def __eq__(self, arg0: Tuple[int, int]) -> bool: ...
    @property
    def ct(self) -> int: ...
    @property
    def rt(self) -> int: ...

def can_use_interactive_placer(policy_type: PolicyType) -> bool: ...
def chip_placement_policy_from_string(chip_placement_policy_str: str) -> ChipPlacementPolicy: ...
def graph_solver_self_cut_type_from_string(graph_solver_self_cut_type_from_string_str: str) -> GraphSolverSelfCutType: ...
def policy_from_string(policy_type_str: str) -> PolicyType: ...
